"""
EnsembleRetriever混合检索：
1、向量检索
2、关键字检索
使用场景：
- 多模态检索需求: 当需要同时使用向量检索(如语义搜索)和关键词检索(如TF-IDF/BM25)时
- 提高召回率: 单一检索方法可能遗漏相关文档，混合检索可以覆盖更多可能性
- 冗余消除: 多个检索器可能返回相似结果，EnsembleRetriever可以去除冗余
"""
from pprint import pprint

from langchain.retrievers import EnsembleRetriever
from langchain_community.document_loaders import TextLoader
from langchain_community.retrievers import BM25Retriever
from langchain_community.vectorstores import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter

from models import get_ollama_embeddings_client

file_path = "../data/document/deepseek百度百科.txt"
docs = TextLoader(file_path, encoding="utf-8").load()

splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=50)
split_docs = splitter.split_documents(docs)

question = "相关评价"

# 向量检索
embeddings_client = get_ollama_embeddings_client()
vector_store = Chroma.from_documents(split_docs, embedding=embeddings_client)
retriever = vector_store.as_retriever(search_kwargs={"k": 3})
retrieve_docs = retriever.invoke(question)
print("-------------------向量检索-------------------------")
pprint(retrieve_docs)

# 关键词检索
bm25_retriever = BM25Retriever.from_documents(split_docs)
BM25Retriever.k = 3
bm25_docs = bm25_retriever.invoke(question)
print("-------------------BM25检索-------------------------")
pprint(bm25_docs)

# 混合检索
ensemble_retriever = EnsembleRetriever(retrievers=[retriever, bm25_retriever], weights=[0.5, 0.5])
ensemble_docs = ensemble_retriever.invoke(question)
print("-------------------混合检索-------------------------")
pprint(ensemble_docs)